import math

import torch
from torch.utils.data import Dataset, DataLoader
from torch.nn import Linear, MSELoss, Sequential, Sigmoid
import torch.optim as optim

import torch
import torch.nn as nn
import numpy as np

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer('pe', pe)

    def forward(self, x):
        '''
        x: [seq_len, batch_size, d_model]
        '''
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)

class SimpleGPT(nn.Module):
    def __init__(self, vocab_size, d_model=256, nhead=8, num_decoder_layers=3):
        super(SimpleGPT, self).__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.positional_encoding = PositionalEncoding(d_model)
        decoder_layer = nn.TransformerDecoderLayer(d_model, nhead)
        self.transformer = nn.TransformerDecoder(decoder_layer, num_decoder_layers)
        self.fc = nn.Linear(d_model, vocab_size)

    def forward(self, src, tgt):
        d_model = 256
        src = self.embedding(src) * np.sqrt(d_model)
        src = self.positional_encoding(src)

        tgt = self.embedding(tgt) * np.sqrt(d_model)
        tgt = self.positional_encoding(tgt)

        memory = self.transformer(src, tgt)
        output = self.fc(memory)
        return output

# 定义训练数据
training_data = [
    ("你因错过太阳而流泪", "你也将错过群星"),
    ("夜幕降临", "群星闪烁"),
    # 可以继续添加更多的训练数据
]


# 数据加载步骤
class MyDataset(Dataset):
    def __init__(self, data):
        self.data = data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]


dataset = MyDataset(training_data)
dataloader = DataLoader(dataset, batch_size=1, shuffle=True)

# 定义模型
model = SimpleGPT(vocab_size)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练步骤
num_epochs = 10
for epoch in range(num_epochs):
    for i, (input_text, target_text) in enumerate(dataloader):
        optimizer.zero_grad()
        input_tensor = torch.Tensor(
            [len(input_text[0]), len(input_text[1]), len(target_text[0]), len(target_text[1]), len(target_text[1])])
        target_tensor = torch.Tensor(
            [len(target_text[0]), len(target_text[1]), len(target_text[0]), len(target_text[1]), len(target_text[1])])

        output = model(input_tensor)
        loss = criterion(output, target_tensor)
        loss.backward()
        optimizer.step()

        if i % 100 == 0:
            print(f'Epoch [{epoch + 1}/{num_epochs}], Step [{i + 1}/{len(dataloader)}], Loss: {loss.item()}')

print("训练完成！")




# 智能问答系统类
class QASystem:
    def __init__(self, model, vocab, max_length=100):
        self.model = model
        self.vocab = vocab
        self.max_length = max_length

    def generate_answer(self, question):
        self.model.eval()
        question_tokens = [self.vocab[word] for word in question.split()]
        question_tensor = torch.LongTensor(question_tokens).unsqueeze(1)

        # 使用起始符号来初始化回答
        answer_input = torch.LongTensor([self.vocab['&ltstart&gt']]).unsqueeze(1)

        for i in range(self.max_length):
            output = self.model(question_tensor, answer_input)
            output_token = output.argmax(dim=2)[-1].item()
            if output_token == self.vocab['&ltend&gt']:
                break
            answer_input = torch.cat((answer_input, torch.LongTensor([[output_token]])), dim=0)

        answer_tokens = [token.item() for token in answer_input]
        answer_words = [word for word, idx in self.vocab.items() if idx in answer_tokens]
        answer = ' '.join(answer_words)
        return answer

# 假设你已经有了一个训练好的模型和词汇表
# vocab是一个字典，将词映射到索引
# model是你训练好的SimpleGPT模型

# 初始化问答系统
qa_system = QASystem(model, vocab)

# 提出问题
question = "为什么错过太阳会流泪？"

# 获取回答
answer = qa_system.generate_answer(question)
print(answer)






